#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

#define HIST_SIZE 256

static cv::Mat _cal_histogram(const cv::Mat& src) {
    CV_Assert(src.type() == CV_8UC1);
    
    /*int hist[HIST_SIZE];
    memset(hist, 0, HIST_SIZE * sizeof(int));
    for (int i = 0; i < src.rows; i++) {
        for (int j = 0; j < src.cols; j++) {
            int idx = (int)src.ptr<uchar>(i)[j];
            if (idx < HIST_SIZE) {
                hist[idx]++;
            }
        }
    }
    cv::Mat histMat(HIST_SIZE, 1, CV_32FC1);
    for (int i = 0; i < HIST_SIZE; i++) {
        histMat.ptr<float>(i)[0] = hist[i];
    }
    return histMat;*/
    int nimages = 1;
    int channels[] = {0};
    int dims = 1;
    int histsizes[] = {HIST_SIZE};
    float histrange[] = {0, HIST_SIZE};
    const float* ranges[] = {histrange};
    bool uniform = true;
    bool accumulate = false;
    cv::Mat histMat;
    cv::calcHist(&src, nimages, channels, cv::Mat(), histMat, dims, histsizes, ranges, uniform, accumulate);
    return histMat;
}

static void _show_histogram(const cv::Mat& mat, cv::String winName) {
    cv::Mat hist = _cal_histogram(mat);
    cv::normalize(hist, hist, 0, 255, cv::NORM_MINMAX, CV_8UC1);
    cv::Mat hist2 = cv::Mat(300, 256, CV_8UC1, cv::Scalar(0));
    for (int i = 0; i < HIST_SIZE; i++) {
        cv::line(hist2, cv::Point(i, 299), cv::Point(i, 299 - hist.at<uchar>(i)), cv::Scalar(255), 1, cv::LINE_AA);
    }
    cv::imshow(winName, hist2);
}

static cv::Mat _histogram_equalize(const cv::Mat& src) {
    CV_Assert(src.type() == CV_8UC1);
    cv::Mat hist = _cal_histogram(src);
    /*
    int hist_accumulate[HIST_SIZE];
    int sum = 0;
    int total_number_pixels = src.rows*src.cols;
    for (int i = 0; i < HIST_SIZE; i++) {
        sum += (int)hist.ptr<float>(i)[0];
        hist_accumulate[i] = cv::saturate_cast<uchar>(sum * 255 / total_number_pixels);
    }
    cv::Mat dst(src.size(), CV_8UC1);
    for (int i = 0; i < dst.rows; i++) {
        for (int j = 0; j < dst.cols; j++) {
            dst.ptr<uchar>(i)[j] = hist_accumulate[src.ptr<uchar>(i)[j]];
        }
    }
    return dst;*/
    
    
    int first_non_zero_hist_pixel_value = 0;
    for (int i = 0; i < hist.rows; i++) {
        int value = (int)hist.ptr<float>(i)[0];
        if (value == 0) {
            first_non_zero_hist_pixel_value++;
        } else {
            break;
        }
    }
    
    cv::Mat dst = cv::Mat(src.size(), src.type());
    int first_non_zero_hist_value = (int)hist.ptr<float>(first_non_zero_hist_pixel_value)[0];
    int total_number_pixels = src.rows * src.cols;
    if (first_non_zero_hist_value == total_number_pixels) {
        dst.setTo(first_non_zero_hist_pixel_value);
        return dst;
    }
    
    float scale = (HIST_SIZE-1.f) / (total_number_pixels - hist.ptr<float>(first_non_zero_hist_pixel_value)[0]);
    int sum = 0;
    int hist_accumulate[HIST_SIZE];
    hist_accumulate[first_non_zero_hist_pixel_value] = 0;
    int pixel_value = first_non_zero_hist_pixel_value+1;
    for (; pixel_value < HIST_SIZE; pixel_value++) {
        sum += (int)hist.ptr<float>(pixel_value)[0];
        hist_accumulate[pixel_value] = cv::saturate_cast<uchar>(sum * scale);
    }
    
    for (int i = 0; i < dst.rows; i++) {
        for (int j = 0; j < dst.cols; j++) {
            dst.ptr<uchar>(i)[j] = hist_accumulate[src.ptr<uchar>(i)[j]];
        }
    }
    return dst;
}

int main(int argc, char **argv) {
    cv::Mat src = cv::imread("/home/xlll/Downloads/opencv/samples/data/chicky_512.png", cv::IMREAD_GRAYSCALE);
    if (src.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat dst;
    cv::equalizeHist(src, dst);
    
    cv::Mat dst2 = _histogram_equalize(src);
    cv::Mat diff = dst != dst2;
    std::vector<cv::Point> nonzeros;
    cv::findNonZero(diff, nonzeros);
    std::cout << "nonzeros size = " << nonzeros.size() << std::endl;
    
    cv::imshow("src", src);
    cv::imshow("dst", dst);
    cv::imshow("dst2", dst2);
    _show_histogram(src, "src_hist");
    _show_histogram(dst, "dst_hist");
    _show_histogram(dst2, "dst2_hist");
    
    cv::waitKey(0);
    
    return 0;
}
